import torch
import cv2
import os
import torch.nn.functional as F
from tqdm import tqdm
import numpy as np
from sklearn.metrics import f1_score
from model import model_resnext50_32x4d,model_Xception,model_InceptionV4
from process_data import get_data_list,concat_data,normalize_data,adjust_data_list
from dataset import Mydataset
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor,Normalize
from collections import Counter

def pred_fun(model,image):
    with torch.no_grad():
        model.eval()
        pred = model(image)
        pred = torch.argmax(F.log_softmax(pred, dim=1), dim=1)
        pred = pred.cpu().numpy().tolist()
        return pred


def val_fun(model,data_loader,stage,show_counter = False):
    pred_list = []
    label_list = []

    with torch.no_grad():
        model.eval()
        for index,(batch_image,batch_label) in enumerate(tqdm(data_loader)):
            ## 返回一个batch 的结果，存在list中。
            batch_image = batch_image.float().cuda()
            pred = pred_fun(model,batch_image)
            pred_list += pred
            label_list += batch_label.numpy().tolist()
    f1 = f1_score(label_list,pred_list,average=None)
    print(f1)
    if stage == 'stage1':
        f1 = 0.25 * f1[0] + 0.3 * f1[1] + 0.45 * f1[2]
    elif stage == 'stage2':
        f1 = 0.4 * f1[0] + 0.6 * f1[1]
    else:
        f1 = 0.1 * f1[0] + 0.2 * f1[1] + 0.3 * f1[2] + 0.4 * f1[3]

    result = list(zip(pred_list, label_list))

    if show_counter:
        ## 将label为1的数据预测结果存入
        pred_0 = [p[0] for p in result if p[1] == 0]
        pred_1 = [p[0] for p in result if p[1] == 1]
        pred_2 = [p[0] for p in result if p[1] == 2]
        pred_3 = [p[0] for p in result if p[1] == 3]
        print(Counter(pred_0))
        print(Counter(pred_1))
        print(Counter(pred_2))
        print(Counter(pred_3))
        print(Counter(label_list))

    return f1


if __name__ == '__main__':
    '''
    val the model
    '''
    Binary = False
    # Image_Size = (800, 488)
    # Image_Size = (800, 450)
    Image_Size = (300, 170)
    # Image_Size = (512, 288)
    Normalize = True
    Concat = 0


    # pre_trained_dir = '../model_save/final/resnext50_concat_diff_lr_all_300/best_.pth'
    # pre_trained_dir = '../model_save/final/0924/resnext50_shuffle_concat_freeze_all_300/best_.pth'
    # pre_trained_dir = '../model_save/final/resnext50_concat_all_300/best_.pth'
    pre_trained_dir = '../model_save/final/1002\InceptionV4_concat0_weight_300/best_.pth'

    # model = model_resnext50_32x4d(Num_Class=4)
    model = model_InceptionV4(num_class=4)
    # model = model_Xception(num_class=3)
    model.load_state_dict(torch.load(pre_trained_dir))
    ## 图片进行concat操作时候，必须要only_key_frame。单独验证时不进行over_sample操作。
    data_list_train,data_list_val = get_data_list(only_key_frame=True,over_sample=False)
    data_list_val = adjust_data_list(data_list_val)
    val_dataSet = Mydataset(data_list_val, train=False, image_size=Image_Size, concat=Concat, normalize=Normalize)
    val_data_loader = DataLoader(dataset=val_dataSet, batch_size=4, shuffle=True,
                                 num_workers=6, pin_memory=True)

    val_acc = val_fun(model,val_data_loader,None,show_counter=True)
    # print('train acc--',train_acc)
    print('--val acc--',val_acc)